package dd.lo.flann;

import dd.lo.HelloCV;
import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.BorderPane;
import javafx.scene.layout.FlowPane;
import javafx.scene.text.Text;
import javafx.stage.FileChooser;
import javafx.stage.Stage;
import org.opencv.core.*;
import org.opencv.features2d.Features2d;
import org.opencv.features2d.FlannBasedMatcher;
import org.opencv.features2d.ORB;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import org.opencv.xfeatures2d.SURF;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

public class FlannMatchExample extends Application {

    public static void main(String[] args) {
        System.load(HelloCV.class.getResource("/libopencv_java481.dylib").getPath());
        FlannMatchExample.start(args);
    }
    public static void start(String[] args) {
        launch(args);
    }

    private static final int DEFAULT_WIDTH = 1080;
    private static final int DEFAULT_HEIGHT = 768;

    private ImageView imageView;

    private Mat img1;
    private Mat img2;
    private Text result;

    private void compareImages() {
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                //Resize img1 to fit the size of img2
                ImageUtils.fitImgSize(img1, img2.width(), img2.height());
                Mat gray1 = new Mat(), gray2 = new Mat();
                Imgproc.cvtColor(img1, gray1, Imgproc.COLOR_BGR2GRAY);
                Imgproc.cvtColor(img2, gray2, Imgproc.COLOR_BGR2GRAY);
                //锐化图像
                Mat kernel = new Mat(3, 3, CvType.CV_32SC1);
//                kernel.put(0, 0, new int[]{0, -1,  0, -1,  5, -1, 0, -1,  0});
                kernel.put(0, 0, new int[]{-1, -1,  -1, -1,  9, -1, -1, -1,  -1});
                Imgproc.filter2D(gray1, gray1, gray1.depth(), kernel);
                Imgproc.filter2D(gray2, gray2, gray2.depth(), kernel);
//                SURF detector = SURF.create(); //Patented, set OPENCV_ENABLE_NONFREE CMake option and rebuild the library in function 'create'
                ORB detector = ORB.create();
                MatOfKeyPoint kp1 = new MatOfKeyPoint();
                //Descriptor must be in type of CV_32F. Ref: https://answers.opencv.org/question/11209/unsupported-format-or-combination-of-formats-in-buildindex-using-flann-algorithm/
                Mat des1 = new Mat();
                detector.detectAndCompute(gray1, new Mat(), kp1, des1);
                des1.convertTo(des1, CvType.CV_32F);
                MatOfKeyPoint kp2 = new MatOfKeyPoint();
                Mat des2 = new Mat();
                detector.detectAndCompute(gray2, new Mat(), kp2, des2);
                des2.convertTo(des2, CvType.CV_32F);
                FlannBasedMatcher flannBasedMatcher = FlannBasedMatcher.create();
                flannBasedMatcher.add(Collections.singletonList(des2));
                List<MatOfDMatch> matches = new ArrayList<>();
                flannBasedMatcher.knnMatch(des1, matches, 2);
                //Filter good matches
                List<MatOfDMatch> goodMatches = new ArrayList<>(matches.size());
                for (MatOfDMatch match : matches) {
                    DMatch[] dMatches = match.toArray();
                    DMatch p1 = dMatches[0], p2 = dMatches[1];
                    if (p1.distance < 0.7 * p2.distance) {
                        goodMatches.add(match);
                    }
                }
                Mat dst = new Mat();
                Features2d.drawMatchesKnn(img1, kp1, img2, kp2, goodMatches, dst);
//                Features2d.drawMatches(gray1, kp1, gray2, kp2, goodMatches, dst, 1, new Scalar(255, 0, 0),  new Scalar(0, 0, 0), new MatOfByte(), Features2d.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS);
                //调整图片大小以用于展示
                ImageUtils.fitImgSize(dst, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                MatOfByte buffer = new MatOfByte();
                Imgcodecs.imencode(".jpeg", dst, buffer);
                Platform.runLater(() -> {
                    imageView.setImage(new Image(new ByteArrayInputStream(buffer.toArray())));
                    result.setText(String.format("发现匹配特征点%d个", goodMatches.size()));
                });
                return null;
            }
        };
        new Thread(task).start();
    }

    @Override
    public void start(Stage stage) {
        Text msg1 = new Text();
        Text msg2 = new Text();
        result = new Text();
        Button imageChooser1 = new Button("选择待比较图片");
        imageChooser1.setOnMouseClicked(e -> {
            File imgFile = chooseFile(stage);
            if (imgFile == null) return;
            //读取一张图片
            img1 = Imgcodecs.imread(imgFile.getAbsolutePath(), Imgcodecs.IMREAD_UNCHANGED);
            if (img1.empty()) {
                throw new RuntimeException("图片读取失败");
            }
            msg1.setText("待比较图片已准备好");
        });
        Button imageChooser2 = new Button("选择基准图片");
        imageChooser2.setOnMouseClicked(e -> {
            File imgFile = chooseFile(stage);
            if (imgFile == null) return;
            //读取一张图片
            img2 = Imgcodecs.imread(imgFile.getAbsolutePath(), Imgcodecs.IMREAD_UNCHANGED);
            if (img2.empty()) {
                throw new RuntimeException("图片读取失败");
            }
            msg2.setText("基准图片已准备好");
        });
        Button compareBtn = new Button("开始比较");
        compareBtn.setOnMouseClicked(event -> {
            if (null == img1 || img1.empty() || null == img2 || img2.empty()) {
                System.out.println("没有准备好2个图片，不能开始");
                return;
            }
            compareImages();
        });
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(imageChooser1, msg1, imageChooser2, msg2, compareBtn, result);
        bottomControl.setPrefHeight(50);
        imageView = new ImageView();
        BorderPane root = new BorderPane();
        root.setCenter(imageView);
        root.setBottom(bottomControl);
        Scene scene = new Scene(root, DEFAULT_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("FLANN matcher example");
        stage.setResizable(false);
        stage.show();
    }

    private static File chooseFile(Stage stage) {
        FileChooser fileChooser = new FileChooser();
        fileChooser.setTitle("请选择JPEG图片");
        fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
        );
        fileChooser.getExtensionFilters().addAll(
                new FileChooser.ExtensionFilter("JPEG", "*.jpeg"),
                new FileChooser.ExtensionFilter("PNG", "*.png")
        );
        File pdfFile = fileChooser.showOpenDialog(stage);
        if (null == pdfFile) {
            System.out.println("用户取消选中文件");
            return null;
        }
        return pdfFile;
    }
}
